module deform_flow_test_mod

  !Kent et al 2014, Dynamical core model intercomparison project: Tracer transport test cases.
  ! DCMIP1-1

	use const_mod
	use namelist_mod
	use sphere_geometry_mod
	use mesh_mod
	use vert_coord_mod
	use state_mod
	use tracer_mod
	use parallel_mod

	implicit none

	private
	public deform_flow_test_init
	public deform_flow_test_set_ic
	public deform_flow_test_set_wind

	real(r8), parameter :: period = 12 * 86400           ! period of motion (12 days)
	real(r8), parameter :: omega0 = 23000 * pi / period  ! maximum of the vertical pressure velocity in units of Pa s-1
	real(r8), parameter :: b      = 0.2                  ! normalized pressure depth of the divergence layer
	real(r8), parameter :: lon1   = 5 * pi / 6					 ! initial longitude of first tracer
	real(r8), parameter :: lon2   = 7 * pi / 6					 ! initial longitude of second tracer
	real(r8), parameter :: latc   = 0					           ! initial latitude of tracers
	real(r8), parameter :: zc     = 5000					       ! initial altitude of tracers (m)
	real(r8), parameter :: Rt     = radius / 2           ! horizontal half-width of tracers
	real(r8), parameter :: zt     = 1000                 ! vertical half-width of tracers
	real(r8), parameter :: T0     = 300                  ! isothermal atmospheric temperature
	real(r8)  ztop, p0
contains

	subroutine deform_flow_test_init()

		ztop = 12000._r8
		p0   = 100000._r8

	end subroutine deform_flow_test_init

	subroutine deform_flow_test_set_ic(mesh, tracer)

		type(mesh_type), intent(in) :: mesh
		type(tracer_type), intent(inout) :: tracer
		integer i, j, k, itracer
		real(r8) lon, lat, r1, r2, d1, d2, p, z

		if (ntracers /= 3) then
			print*, 'ntracers must be 3!'
			stop
		end if
 
    do k = mesh%full_lev_ibeg, mesh%full_lev_iend
    	p = hybrid_coord_calc_ph(k, p0)
    	z = Rd * T0 / g * log(p0 / p)
    	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    		lat = mesh%full_lat(j)
    		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    			lon = mesh%full_lon(i)
    			r1 = calc_distance(lon, lat, lon1, latc)
    			r2 = calc_distance(lon, lat, lon2, latc)
    			d1 = min(1.0, (r1 / Rt)**2 + ((z - zc) / zt)**2)
    			d2 = min(1.0, (r2 / Rt)**2 + ((z - zc) / zt)**2)
    			tracer%q(i,j,k,1) = 0.5 * ((1 + cos(pi * d1)) + (1 + cos(pi * d2)))
    			tracer%q(i,j,k,2) = 0.9 - 0.8 * tracer%q(i,j,k,1)**2
    			if (d1 < 0.5 ) then
    				tracer%q(i,j,k,3) = 1
    			else if (d2 < 0.5) then
    				tracer%q(i,j,k,3) = 1
    			else
    				tracer%q(i,j,k,3) = 0.1
    			end if
    			if (z > zc .and. (lat > latc - 1./8 .and. lat < latc + 1./8)) then
    				tracer%q(i,j,k,3) = 0.1
    			end if
    		end do
    	end do
    end do

    ! do k = mesh%full_lev_ibeg, mesh%full_lev_iend
    ! 	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    ! 		do i = mesh%full_lon_ibeg, mesh%full_lon_iend	
    ! 			tracer%q(i,j,k,4) = 1 - 0.3 * (tracer%q(i,j,k,1) + tracer%q(i,j,k,2) + tracer%q(i,j,k,3))
    ! 		end do
    ! 	end do
    ! end do

    do itracer = 1, size(tracer%q, 4)
    	call fill_zonal_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    	call fill_merid_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    end do

	end subroutine deform_flow_test_set_ic

	subroutine deform_flow_test_set_wind(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds
		real(r8) lon, lat, p, ptop, cos_lat, cos_t, sp, u0, w_scale
		integer i, j, k

		w_scale = 4.0_r8

		ptop = p0 * exp(-g * ztop / Rd / T0)
		cos_t = cos(pi * time_in_seconds / period)
		u0 = 10 * radius / period

		associate (u => state%u, v => state%v, we => state%we, phs => state%phs)
		phs = p0
		do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			do j = mesh%full_lat_ibeg+1, mesh%full_lat_iend-1
				lat = mesh%full_lat(j)
				do i = mesh%half_lon_ibeg, mesh%half_lon_iend
					lon = mesh%half_lon(i) - pi2 * time_in_seconds / period
					p = hybrid_coord_calc_ph(k, phs(i,j))
					u(i,j,k) = u0 * sin(lon)**2 * sin(2 * lat) * cos_t + pi2 * radius / period * cos(lat) + &
						         omega0 * radius / b / ptop * cos(lon) * cos(lat)**2 * cos(pi2 * time_in_seconds / period) * &
						         (-exp((p - p0) / (b * ptop)) + exp((ptop - p) / (b * ptop))) * w_scale
				end do
			end do
		end do

		do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			do j = mesh%half_lat_ibeg, mesh%half_lat_iend
				cos_lat = mesh%half_cos_lat(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					lon = mesh%full_lon(i) - pi2 * time_in_seconds / period
					v(i,j,k) = u0 * sin(2 * lon) * cos_lat * cos_t
				end do
			end do
		end do

	  call fill_zonal_halo_lon(u, all_halo=.true.)

	  do k = mesh%half_lev_ibeg, mesh%half_lev_iend
	  	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
	  		cos_lat = mesh%full_cos_lat(j)
	  		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
	  			lon = mesh%full_lon(i) - pi2 * time_in_seconds / period
	  			p = hybrid_coord_calc_ph_lev(k, phs(i,j))
	  			sp = 1 + exp((ptop - p0) / b / ptop) - exp((p - p0) / b / ptop) - exp((ptop - p) / b / ptop)
	  			we(i,j,k) = omega0 * sin(lon) * cos_lat * cos(pi2 * time_in_seconds / period) * sp * w_scale
	  		end do
	  	end do
	  end do
	  
	end associate

	end subroutine deform_flow_test_set_wind

end module deform_flow_test_mod